# 三维扩散方程
# Gitee Repo

import numpy as np
import scipy
import torch

L  = 2;
dx = 0.05;
dt = 0.0001;
D  = 1

torch.set_default_device('cuda' if torch.cuda.is_available() else 'cpu')

x,y,z = torch.meshgrid(torch.arange(-L,L,dx),torch.arange(-L,L,dx),torch.arange(-L,L,dx),indexing='ij')
n = x.shape[0]
sigma = dt*D/(dx**2)

u0 = torch.zeros((n,n,n));
u1 = torch.zeros((n,n,n));

u0[(x>-0.2*L) & (x<0.2*L) & (y>-0.2*L) & (y<0.2*L) & (z>-0.2*L) & (z<0.2*L)]=1;

print(u0.device)

for tick in range(1000):
    print(tick)
    diff_i = u0[2:n,1:n-1,1:n-1]-2*u0[1:n-1,1:n-1,1:n-1]+u0[0:n-2,1:n-1,1:n-1]
    diff_j = u0[1:n-1,2:n,1:n-1]-2*u0[1:n-1,1:n-1,1:n-1]+u0[1:n-1,0:n-2,1:n-1]
    diff_k = u0[1:n-1,1:n-1,2:n]-2*u0[1:n-1,1:n-1,1:n-1]+u0[1:n-1,1:n-1,0:n-2]
    
    u1[1:n-1,1:n-1,1:n-1] = u0[1:n-1,1:n-1,1:n-1]+sigma*(diff_i+diff_j+diff_k)
    
    u0,u1=u1,u0;

u_cpu = u1.cpu().numpy()[::2,::2,::2]
x_cpu = x.cpu().numpy()[::2,::2,::2]
y_cpu = y.cpu().numpy()[::2,::2,::2]
z_cpu = z.cpu().numpy()[::2,::2,::2]

scipy.io.savemat('diff3D_data.mat',{'x':x_cpu,'y':y_cpu,'z':z_cpu,'u':u_cpu})
exit()
